from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import six
import os
import os.path as osp
import sys
import copy
import logging
from ast import literal_eval

import numpy as np
import yaml

__all__ = ['cfg', 'cfg_from_file', 'load_cfg_file']

class AttrDict(dict):

    IMMUTABLE = '__immutable__'

    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__[AttrDict.IMMUTABLE] = False

    def __getattr__(self, name):
        if name in self.__dict__:
            return self.__dict__[name]
        elif name in self:
            return self[name]
        else:
            raise AttributeError(name)

    def __setattr__(self, name, value):
        if not self.__dict__[AttrDict.IMMUTABLE]:
            if name in self.__dict__:
                self.__dict__[name] = value
            else:
                self[name] = value
        else:
            raise AttributeError(
                'Attempted to set "{}" to "{}", but AttrDict is immutable'.
                format(name, value)
            )

    def immutable(self, is_immutable):
        """Set immutability to is_immutable and recursively apply the setting
        to all nested AttrDicts.
        """
        self.__dict__[AttrDict.IMMUTABLE] = is_immutable
        # Recursively set immutable state
        for v in self.__dict__.values():
            if isinstance(v, AttrDict):
                v.immutable(is_immutable)
        for v in self.values():
            if isinstance(v, AttrDict):
                v.immutable(is_immutable)

    def is_immutable(self):
        return self.__dict__[AttrDict.IMMUTABLE]


__C = AttrDict()
# Consumers can get config by:
#   from fast_rcnn_config import cfg
cfg = __C


# Random note: avoid using '.ON' as a config key since yaml converts it to True;
# prefer 'ENABLED' instead

# ---------------------------------------------------------------------------- #
# Dicom options
# ---------------------------------------------------------------------------- #
__C.DICOM = AttrDict()

# Dicom fliter sequnce description
__C.DICOM.BODY_TAGS = ("CHEST",)

# Diocom fliter flag
__C.DICOM.DESC_FLITER_GLAG = True

# DICOM protocolName tag
__C.DICOM.PROTOCOL_TAG = (0x0018, 0x1030)

# Dicom sliceThinkness tag
__C.DICOM.THINKNESS_TAG = (0x0018, 0x0050)

# Dicom series description
__C.DICOM.SERIES_DESC_TAG = (0x0008, 0x103e)

# Dicom image uid
__C.DICOM.SLICE_IMAGE_UID_NAME = "MediaStorageSOPInstanceUID"


# ---------------------------------------------------------------------------- #
# Data loader options
# ---------------------------------------------------------------------------- #
__C.DATA_LOADER = AttrDict()

# Number of Python threads to use for the data loader (warning: using too many
# threads can cause GIL-based interference with Python Ops leading to *slower*
# training; 4 seems to be the sweet spot in our experience)
__C.DATA_LOADER.TRAIN_NUM_WORKERS = 8

__C.DATA_LOADER.VALID_NUM_WORKERS = 8

__C.DATA_LOADER.TEST_NUM_WORKERS = 8

__C.DATA_LOADER.INFERENCE_NUM_WORKERS = 4

#Batch size
__C.DATA_LOADER.BS = 2

__C.DATA_LOADER.CROP = False

__C.DATA_LOADER.VAL_PERIOD = 10

# ---------------------------------------------------------------------------- #
#Optimizer
# ---------------------------------------------------------------------------- #
__C.SOLVER = AttrDict()

__C.SOLVER.OPTIMIZER = 'SGD'

__C.SOLVER.OPTIMIZER_PARAMETER = {}

__C.SOLVER.SCHEDULER = 'StepLR'

__C.SOLVER.SCHEDULER_PARAMETER = {}

__C.SOLVER.SCHEDULER_WARMUP = {}

__C.SOLVER.LOSS = 'ElasticLoss'

__C.SOLVER.MAX_STEPS = 10

# ---------------------------------------------------------------------------- #
# Model options
# ---------------------------------------------------------------------------- #
__C.MODEL = AttrDict()

# Stats
__C.MODEL.STATS = [96.70380979, 58.26832049]
# Scales
__C.MODEL.SCALES = (512, 512)


# ---------------------------------------------------------------------------- #
# Keypoint Regression options
# ---------------------------------------------------------------------------- #
__C.MODEL.LANDMARK_DET = AttrDict()

__C.MODEL.LANDMARK_DET.DEVICE = 'cpu'

__C.MODEL.LANDMARK_DET.CHANNELS = 3

__C.MODEL.LANDMARK_DET.KP_NUM = 68

__C.MODEL.LANDMARK_DET.HEAT_SIZE = (128, 64)

__C.MODEL.LANDMARK_DET.REG_TYPE = 'kp'

__C.MODEL.LANDMARK_DET.MODEL = 'HourglassNet'

__C.MODEL.LANDMARK_DET.MODEL_PARAMETER = {'bw': 24,
                                                 'hg_depth': 4,
                                                 'upmode':'bilinear',
                                                 'multiscale_hg_block': True,
                                                 'se':False,
                                                 'se_ratio':16,
                                                 'use_drop':False}

__C.MODEL.LANDMARK_DET.SHAPE = (512, 256)

__C.MODEL.LANDMARK_DET.CLASS = {}
__C.MODEL.LANDMARK_DET.THRESH = 0.5

__C.MODEL.LANDMARK_DET.STATS = []

__C.MODEL.LANDMARK_DET.PRETRAINED = ""

__C.MODEL.LANDMARK_DET.RESUME = ""


# ---------------------------------------------------------------------------- #
# Semantic segmantation options
# ---------------------------------------------------------------------------- #
__C.MODEL.SEMANTIC_SEG = AttrDict()
__C.MODEL.SEMANTIC_SEG.DEVICE = 'cpu'
__C.MODEL.SEMANTIC_SEG.CHANNELS = 3
__C.MODEL.SEMANTIC_SEG.NUM_CLASSES = 2
__C.MODEL.SEMANTIC_SEG.MODEL = 'UNet'

__C.MODEL.SEMANTIC_SEG.MODEL_PARAMETER = {}
__C.MODEL.SEMANTIC_SEG.SHAPE = (512, 256)

__C.MODEL.SEMANTIC_SEG.PRE_TRAIN_DIR = ""

__C.MODEL.SEMANTIC_SEG.CLASS = {}

__C.MODEL.SEMANTIC_SEG.THRESH = 0.5

__C.MODEL.SEMANTIC_SEG.PRETRAINED = ''

__C.MODEL.SEMANTIC_SEG.RESUME = ""

__C.MODEL.SEMANTIC_SEG.STATS = []




# ---------------------------------------------------------------------------- #
# Angles Regression options
# ---------------------------------------------------------------------------- #
__C.MODEL.ANGLES_REGRESSION = AttrDict()

__C.MODEL.ANGLES_REGRESSION.CHANNELS = 3

__C.MODEL.ANGLES_REGRESSION.NUM_ANGLES = 3

__C.MODEL.ANGLES_REGRESSION.MODEL = 'HourglassNet'

__C.MODEL.ANGLES_REGRESSION.MODEL_PARAMETER = {'bw': 24,
                                                 'hg_depth': 4,
                                                 'upmode':'bilinear',
                                                 'multiscale_hg_block': True,
                                                 'se':False,
                                                 'se_ratio':16,
                                                 'use_drop':False}

__C.MODEL.ANGLES_REGRESSION.SHAPE = (512, 256)

__C.MODEL.ANGLES_REGRESSION.CLASS = {}

__C.MODEL.ANGLES_REGRESSION.THRESH = 0.5

__C.MODEL.ANGLES_REGRESSION.PRETRAINED = None

__C.MODEL.ANGLES_REGRESSION.RESUME = ""




# ---------------------------------------------------------------------------- #
# Infer options
# ---------------------------------------------------------------------------- #
__C.INFER = AttrDict()
__C.INFER.PRE_PROCESS = AttrDict()
__C.INFER.PRE_PROCESS.SKULL_TEMPLATE = ''


# ---------------------------------------------------------------------------- #
# Dataset options
# ---------------------------------------------------------------------------- #
__C.DATASET = AttrDict()
__C.DATASET.ROOT = ''
__C.DATASET.DATA_TYPE = ''


# ---------------------------------------------------------------------------- #
# Display every options
# ---------------------------------------------------------------------------- #
__C.EVERY = AttrDict()
__C.EVERY.DISPLAY = 10
__C.EVERY.LOG = 10
__C.EVERY.CHECKPOINT = 10
__C.EVERY.VALIDATE = 10

# ---------------------------------------------------------------------------- #
# detection results options
# ---------------------------------------------------------------------------- #
__C.RESULT = AttrDict()

# 0 success
__C.RESULT.SUCCESS = 0

# 1 input path not exist
__C.RESULT.FAILED_NONE_PATH = 1

# 2 access gpu error
__C.RESULT.FAILED_GPU_ERROR = 2

# 3 ct serires description wrong
__C.RESULT.FAILED_DICOM_DESC_ERROR = 3

# 4 preprocess error
__C.RESULT.FAILED_PRRPROCESS_ERROR = 4


# 4 ct sequence length is not satify minimum length
__C.RESULT.FAILED_CT_SEQUENCE_ERROR = 5


# ---------------------------------------------------------------------------- #
# Logging options
# ---------------------------------------------------------------------------- #
__C.LOG = AttrDict()

# Logging level
__C.LOG.LOG_LEVEL = logging.DEBUG

# Logging dir
__C.LOG.DIR = './log'

# ---------------------------------------------------------------------------- #
# MISC options
# ---------------------------------------------------------------------------- #

# Number of GPUs to use (applies to both training and testing)
__C.NUM_GPUS = 1


# For reproducibility
__C.RNG_SEED = 123

# A small number that's used many times
__C.EPS = 1e-14

# Root directory of project
__C.ROOT_DIR = osp.abspath(osp.join(osp.dirname(__file__), '..', '..'))

# Output basedir
__C.OUTPUT_DIR = 'Outputs'

# Dump detection visualizations
__C.VIS = False

# [Infered value]
__C.CUDA = True

__C.DEBUG = False

__C.OUTPUT_DIR = ''



def merge_cfg_from_file(cfg_filename):
    """Load a yaml config file and merge it into the global config."""
    with open(cfg_filename, 'r') as f:
        yaml_cfg = AttrDict(yaml.load(f, Loader=yaml.SafeLoader))
    _merge_a_into_b(yaml_cfg, __C)

cfg_from_file = merge_cfg_from_file


def merge_cfg_from_cfg(cfg_other):
    """Merge `cfg_other` into the global config."""
    _merge_a_into_b(cfg_other, __C)


def merge_cfg_from_list(cfg_list):
    """Merge config keys, values in a list (e.g., from command line) into the
    global config. For example, `cfg_list = ['TEST.NMS', 0.5]`.
    """
    assert len(cfg_list) % 2 == 0
    for full_key, v in zip(cfg_list[0::2], cfg_list[1::2]):
        # if _key_is_deprecated(full_key):
        #     continue
        # if _key_is_renamed(full_key):
        #     _raise_key_rename_error(full_key)
        key_list = full_key.split('.')
        d = __C
        for subkey in key_list[:-1]:
            assert subkey in d, 'Non-existent key: {}'.format(full_key)
            d = d[subkey]
        subkey = key_list[-1]
        assert subkey in d, 'Non-existent key: {}'.format(full_key)
        value = _decode_cfg_value(v)
        value = _check_and_coerce_cfg_value_type(
            value, d[subkey], subkey, full_key
        )
        d[subkey] = value

cfg_from_list = merge_cfg_from_list


def _merge_a_into_b(a, b, stack=None):
    """Merge config dictionary a into config dictionary b, clobbering the
    options in b whenever they are also specified in a.
    """
    assert isinstance(a, AttrDict), 'Argument `a` must be an AttrDict'
    assert isinstance(b, AttrDict), 'Argument `b` must be an AttrDict'

    for k, v_ in a.items():
        full_key = '.'.join(stack) + '.' + k if stack is not None else k
        # a must specify keys that are in b
        if k not in b:
            # if _key_is_deprecated(full_key):
            #     continue
            # elif _key_is_renamed(full_key):
            #     _raise_key_rename_error(full_key)
            # else:
            raise KeyError('Non-existent config key: {}'.format(full_key))

        v = copy.deepcopy(v_)
        v = _decode_cfg_value(v)
        v = _check_and_coerce_cfg_value_type(v, b[k], k, full_key)

        # Recursively merge dicts
        if isinstance(v, AttrDict):
            try:
                stack_push = [k] if stack is None else stack + [k]
                _merge_a_into_b(v, b[k], stack=stack_push)
            except BaseException:
                raise
        else:
            b[k] = v


def _decode_cfg_value(v):
    """Decodes a raw config value (e.g., from a yaml config files or command
    line argument) into a Python object.
    """
    # Configs parsed from raw yaml will contain dictionary keys that need to be
    # converted to AttrDict objects
    if isinstance(v, dict):
        return AttrDict(v)
    # All remaining processing is only applied to strings
    if not isinstance(v, six.string_types):
        return v
    # Try to interpret `v` as a:
    #   string, number, tuple, list, dict, boolean, or None
    try:
        v = literal_eval(v)
    # The following two excepts allow v to pass through when it represents a
    # string.
    #
    # Longer explanation:
    # The type of v is always a string (before calling literal_eval), but
    # sometimes it *represents* a string and other times a data structure, like
    # a list. In the case that v represents a string, what we got back from the
    # yaml parser is 'foo' *without quotes* (so, not '"foo"'). literal_eval is
    # ok with '"foo"', but will raise a ValueError if given 'foo'. In other
    # cases, like paths (v = 'foo/bar' and not v = '"foo/bar"'), literal_eval
    # will raise a SyntaxError.
    except ValueError:
        pass
    except SyntaxError:
        pass
    return v


def _check_and_coerce_cfg_value_type(value_a, value_b, key, full_key):
    """Checks that `value_a`, which is intended to replace `value_b` is of the
    right type. The type is correct if it matches exactly or is one of a few
    cases in which the type can be easily coerced.
    """
    # The types must match (with some exceptions)
    type_b = type(value_b)
    type_a = type(value_a)
    if type_a is type_b:
        return value_a

    # Exceptions: numpy arrays, strings, tuple<->list
    if isinstance(value_b, np.ndarray):
        value_a = np.array(value_a, dtype=value_b.dtype)
    elif isinstance(value_b, six.string_types):
        value_a = str(value_a)
    elif isinstance(value_a, tuple) and isinstance(value_b, list):
        value_a = list(value_a)
    elif isinstance(value_a, list) and isinstance(value_b, tuple):
        value_a = tuple(value_a)
    elif isinstance(value_a, AttrDict) and isinstance(value_b, dict):
        value_a= dict(value_a)
    else:
        raise ValueError(
            'Type mismatch ({} vs. {}) with values ({} vs. {}) for config '
            'key: {}'.format(type_b, type_a, value_b, value_a, full_key)
        )
    return value_a


def load_cfg_file(cfg_file):
    if cfg_file is not None and os.path.exists(cfg_file):
        cfg_from_file(cfg_file)
        print("load  yaml: {}".format(cfg_file))
    else:
        sys.exit(1)
